package com.xiang.WebSocket;

import java.util.List;
import java.util.Map;

import jakarta.websocket.HandshakeResponse;
import jakarta.websocket.server.HandshakeRequest;
import jakarta.websocket.server.ServerEndpointConfig;
import jakarta.websocket.server.ServerEndpointConfig.Configurator;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.server.standard.ServerEndpointExporter;
import org.springframework.web.socket.server.standard.ServletServerContainerFactoryBean;

@Configuration
@EnableWebSocket
public class WebSocketConfig extends Configurator {
  /**
   * 如果使用Springboot默认内置的tomcat容器，则必须注入ServerEndpoint的bean；
   * 如果使用外置的web容器，则不需要提供ServerEndpointExporter，下面的注入可以注解掉
   */
  @Bean
  public ServerEndpointExporter serverEndpointExporter() {
    return new ServerEndpointExporter();
  }

  @Bean
  public ServletServerContainerFactoryBean createWebSocketContainer() {
    ServletServerContainerFactoryBean container = new ServletServerContainerFactoryBean();
    // 在此处设置bufferSize
    container.setMaxTextMessageBufferSize(51200000);
    container.setMaxBinaryMessageBufferSize(51200000);
    container.setMaxSessionIdleTimeout(15 * 60000L);
    return container;
  }

  public static final String IP = "X-Forwarded-For";

  @Override
  public void modifyHandshake(ServerEndpointConfig sec, HandshakeRequest request, HandshakeResponse response) {
    Map<String, List<String>> headers = request.getHeaders();

    if (headers != null) {
      Map<String, Object> userProperties = sec.getUserProperties();
      List<String> ips = headers.get(IP);
      if (ips != null && ips.size() > 0) {
        userProperties.put(IP, ips.get(0));
      }
    }
  }
}
